from flask import Flask, request, jsonify
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography import x509
from cryptography.x509.oid import NameOID
from datetime import datetime, timedelta
import os

app = Flask(__name__)

class CertificateAuthority:
    def __init__(self):
        # 生成CA密钥对
        self.ca_key = rsa.generate_private_key(
            public_exponent=65537,
            key_size=2048
        )
        # 创建自签名根证书
        self.root_cert = self._create_root_certificate()
        
    def _create_root_certificate(self):
        subject = issuer = x509.Name([
            x509.NameAttribute(NameOID.COUNTRY_NAME, "CN"),
            x509.NameAttribute(NameOID.ORGANIZATION_NAME, "CA Authority"),
            x509.NameAttribute(NameOID.COMMON_NAME, "CA Root"),
        ])
        
        return x509.CertificateBuilder().subject_name(
            subject
        ).issuer_name(
            issuer
        ).public_key(
            self.ca_key.public_key()
        ).serial_number(
            x509.random_serial_number()
        ).not_valid_before(
            datetime.utcnow()
        ).not_valid_after(
            datetime.utcnow() + timedelta(days=365 * 10)  # 10年有效期
        ).add_extension(
            x509.BasicConstraints(ca=True, path_length=None), critical=True
        ).add_extension(
            x509.KeyUsage(
                digital_signature=True,
                content_commitment=True,
                key_encipherment=False,
                data_encipherment=False,
                key_agreement=False,
                key_cert_sign=True,
                crl_sign=True,
                encipher_only=False,
                decipher_only=False
            ), critical=True
        ).sign(self.ca_key, hashes.SHA256())

    def issue_certificate(self, csr_pem, validity_days=365):
        # 加载CSR
        csr = x509.load_pem_x509_csr(csr_pem.encode('utf-8'))
        
        # 验证CSR签名
        if not csr.is_signature_valid:
            raise ValueError("Invalid CSR signature")
        
        # 创建证书
        cert = x509.CertificateBuilder().subject_name(
            csr.subject
        ).issuer_name(
            self.root_cert.subject
        ).public_key(
            csr.public_key()
        ).serial_number(
            x509.random_serial_number()
        ).not_valid_before(
            datetime.utcnow()
        ).not_valid_after(
            datetime.utcnow() + timedelta(days=validity_days)
        )
        
        # 添加扩展
        for ext in csr.extensions:
            cert = cert.add_extension(ext.value, critical=ext.critical)
            
        # 签名证书
        signed_cert = cert.sign(
            private_key=self.ca_key,
            algorithm=hashes.SHA256()
        )
        
        return signed_cert.public_bytes(serialization.Encoding.PEM).decode('utf-8')

# 初始化CA
ca = CertificateAuthority()

@app.route('/ca_root', methods=['GET'])
def get_ca_root():
    """获取CA根证书"""
    root_cert_pem = ca.root_cert.public_bytes(
        serialization.Encoding.PEM
    ).decode('utf-8')
    return root_cert_pem, 200, {'Content-Type': 'application/x-pem-file'}

@app.route('/request_cert', methods=['POST'])
def request_certificate():
    """颁发新证书"""
    data = request.json
    if not data or 'csr' not in data:
        return jsonify({'error': 'Missing CSR in request'}), 400
    
    try:
        # 颁发证书
        certificate = ca.issue_certificate(
            data['csr'],
            validity_days=data.get('validity_days', 365)
        )
        return jsonify({'certificate': certificate}), 200
    except Exception as e:
        return jsonify({'error': str(e)}), 400

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=28290, ssl_context='adhoc')